package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"time"
)

type Request struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
	Stream   bool      `json:"stream"`
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type Response struct {
	Model              string    `json:"model"`
	CreatedAt          time.Time `json:"created_at"`
	Message            Message   `json:"message"`
	Done               bool      `json:"done"`
	TotalDuration      int64     `json:"total_duration"`
	LoadDuration       int       `json:"load_duration"`
	PromptEvalCount    int       `json:"prompt_eval_count"`
	PromptEvalDuration int       `json:"prompt_eval_duration"`
	EvalCount          int       `json:"eval_count"`
	EvalDuration       int64     `json:"eval_duration"`
}

// https://mp.weixin.qq.com/s/RCaxYLfd2d1GTAqiqstdUw
func main() {
	client := &http.Client{}
	reader := bufio.NewReader(os.Stdin)
	// start := time.Now()
	fmt.Println("欢迎使用本地 AI 聊天应用！输入 'exit' 退出。")

	for {
		fmt.Print("你: ")
		input, _ := reader.ReadString('\n')
		input = input[:len(input)-1]

		if input == "exit" {
			fmt.Println("再见！")
			break
		}

		message := Message{
			Role:    "user",
			Content: input,
		}

		requestBody, _ := json.Marshal(Request{
			Model:    "llama3.2",
			Messages: []Message{message},
			Stream:   false,
		})

		req, _ := http.NewRequest("POST", "http://localhost:11434/api/chat", bytes.NewBuffer(requestBody))
		req.Header.Set("Content-Type", "application/json")

		resp, err := client.Do(req)
		if err != nil {
			fmt.Println("请求失败:", err)
			continue
		}
		defer resp.Body.Close()

		response := Response{}
		err = json.NewDecoder(resp.Body).Decode(&response)
		fmt.Println(response.Message.Content)
	}
}
